import os
from openpyxl import Workbook
from openpyxl.styles import Font, Alignment, PatternFill, Border, Side
from openpyxl.utils import get_column_letter

class ExcelTool:
    """
    一个用于创建和操作 Excel 文件的工具类。
    """

    def __init__(self):
        """
        初始化一个新的 Excel 工作簿。
        """
        self.wb = Workbook()
        self.ws = self.wb.active  # 默认激活第一个工作表

    def switch_sheet(self, sheet_name:str):
        """
        切换表格
        :param sheet_name: 表格名称
        """
        if sheet_name in self.wb.sheetnames:
            self.ws = self.wb[sheet_name]
        else:
            self.ws =  self.wb.create_sheet(title=sheet_name)
        self.set_sheet_title(sheet_name)

    def set_sheet_title(self, title: str):
        """
        设置当前工作表的标题。
        :param title: 工作表标题
        """
        self.ws.title = title

    def write_data(self, headers: list, data: list, start_row: int = 1, start_col: int = 1):
        """
        写入表头和数据到工作表。
        :param headers: 表头列名列表，例如 ["姓名", "年龄", "城市"]
        :param data: 数据二维列表，例如 [["张三", 25, "北京"], ["李四", 30, "上海"]]
        :param start_row: 起始行号（默认为 1）
        :param start_col: 起始列号（默认为 1）
        """
        # 写入表头
        for col_idx, header in enumerate(headers, start=start_col):
            self.ws.cell(row=start_row, column=col_idx, value=header)

        # 写入数据
        for row_idx, row_data in enumerate(data, start=start_row + 1):
            for col_idx, value in enumerate(row_data, start=start_col):
                self.ws.cell(row=row_idx, column=col_idx, value=value)

    def get_cell_range(self, num_cols, num_rows, start_row: int = 1, start_col: int = 1):
        """
        根据表头和数据自动确定单元格范围。
        :param num_cols: 列数
        :param num_rows: 行数
        :param start_row: 数据起始行号（默认为 1）
        :param start_col: 数据起始列号（默认为 1）
        :return: 单元格范围字符串，例如 "A1:C4"
        """

        # 计算起始和结束的列字母
        start_col_letter = get_column_letter(start_col)
        end_col_letter = get_column_letter(start_col + num_cols - 1)

        # 计算结束行号
        end_row = start_row + num_rows - 1

        # 构造范围字符串
        cell_range = f"{start_col_letter}{start_row}:{end_col_letter}{end_row}"
        return cell_range

    def apply_styles(self, cell_range: str, font=None, alignment=None, fill=None, border=None):
        """
        对指定范围内的单元格应用样式。
        :param cell_range: 单元格范围，例如 "A1:C3"
        :param font: 字体样式（Font 对象）
        :param alignment: 对齐方式（Alignment 对象）
        :param fill: 填充样式（PatternFill 对象）
        :param border: 边框样式（Border 对象）
        """
        for row in self.ws[cell_range]:
            for cell in row:
                if font:
                    cell.font = font
                if alignment:
                    cell.alignment = alignment
                if fill:
                    cell.fill = fill
                if border:
                    cell.border = border

    def auto_adjust_columns_width(self):
        """
        自动调整所有列的宽度以适应内容。
        """
        for col in range(1, self.ws.max_column + 1):
            max_length = 0
            column_letter = get_column_letter(col)
            for cell in self.ws[column_letter]:
                try:
                    if len(str(cell.value)) > max_length:
                        max_length = len(str(cell.value))
                except TypeError:
                    pass
            adjusted_width = (max_length + 10) if max_length > 0 else 10
            self.ws.column_dimensions[column_letter].width = adjusted_width

    def save(self, file_path: str):
        """
        保存 Excel 文件。
        :param file_path: 文件路径（包括文件名和扩展名）
        """
        os.makedirs(os.path.dirname(file_path), exist_ok=True)  # 确保目录存在
        self.wb.save(file_path)
        print(f"Excel 文件已保存到 {file_path}")

    def close(self):
        """
        关闭工作簿。
        """
        self.wb.close()

# 示例用法
if __name__ == "__main__":
    # 创建工具实例
    excel_tool = ExcelTool()

    # 设置工作表标题
    excel_tool.set_sheet_title("示例工作表")

    # 写入数据
    header = ["姓名", "年龄", "城市"]
    data = [
        ["张三", 25, "北京"],
        ["李四", 30, "上海"],
        ["王五", 28, "广州"]
    ]
    excel_tool.write_data(header, data)

    # 应用样式
    # 表头样式
    header_font = Font(bold=True, color="FFFFFF")
    header_fill = PatternFill(start_color="4F81BD", end_color="4F81BD", fill_type="solid")
    header_alignment = Alignment(horizontal="center", vertical="center")
    border = Border(left=Side(style='thin'), right=Side(style='thin'),
                    top=Side(style='thin'), bottom=Side(style='thin'))
    header_cell_range = excel_tool.get_cell_range(len(header), 1)
    excel_tool.apply_styles(header_cell_range, font=header_font, fill=header_fill, alignment=header_alignment, border=border)

    # 内容样式
    content_font = Font(bold=False, color="000000")
    content_fill = None
    content_alignment = Alignment(horizontal="center", vertical="center")
    content_cell_range = excel_tool.get_cell_range(len(header), len(data), start_row=2)
    excel_tool.apply_styles(content_cell_range, font=content_font, fill=content_fill, alignment=content_alignment, border=None)

    # 自动调整列宽
    excel_tool.auto_adjust_columns_width()

    # 保存文件
    excel_tool.save("./result/example.xlsx")

    # 关闭工作簿
    excel_tool.close()